package org.apache.tinkerpop.gremlin.server.util;

import org.apache.tinkerpop.gremlin.groovy.engine.GremlinExecutor;
import org.apache.tinkerpop.gremlin.jsr223.GremlinScriptEngine;
import org.apache.tinkerpop.gremlin.process.traversal.TraversalSource;
import org.apache.tinkerpop.gremlin.server.GraphManager;
import org.apache.tinkerpop.gremlin.server.Settings;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.SimpleBindings;
import java.lang.reflect.Constructor;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.stream.Collectors;

import static org.studiox.graph.common.GraphConstant.JanusGraphConfigConstant.CONFIG_PREFIX;

public class GraphServerGremlinExecutor {
  private static final Logger logger = LoggerFactory.getLogger(GraphServerGremlinExecutor.class);

  private final GraphManager graphManager;
  private final Settings settings;
  private final List<LifeCycleHook> hooks;

  private final ScheduledExecutorService scheduledExecutorService;
  private final ExecutorService gremlinExecutorService;
  private final GremlinExecutor gremlinExecutor;

  private final Map<String, Object> hostOptions = new ConcurrentHashMap<>();

  public GraphServerGremlinExecutor(
      final Settings settings,
      final ExecutorService gremlinExecutorService,
      final ScheduledExecutorService scheduledExecutorService,
      final Graph graph) {
    this.settings = settings;

    try {
      final Class<?> clazz = Class.forName(settings.graphManager);
      final Constructor c = clazz.getConstructor(Settings.class);
      graphManager = (GraphManager) c.newInstance(settings);
    } catch (ClassNotFoundException e) {
      logger.error(
          "Could not find GraphManager implementation "
              + "defined by the 'graphManager' setting as: {}",
          settings.graphManager);
      throw new RuntimeException(e);
    } catch (Exception e) {
      logger.error(
          "Could not invoke constructor on class {} (defined by "
              + "the 'graphManager' setting) with one argument of "
              + "class Settings",
          settings.graphManager);
      throw new RuntimeException(e);
    }

    if (null != graphManager && null != graph) {
      graphManager.putGraph(CONFIG_PREFIX, graph);
    }

    if (null == gremlinExecutorService) {
      final ThreadFactory threadFactoryGremlin = ThreadFactoryUtil.create("exec-%d");
      final BlockingQueue<Runnable> queue = new ArrayBlockingQueue<>(settings.maxWorkQueueSize);
      this.gremlinExecutorService =
          new ThreadPoolExecutor(
              settings.gremlinPool,
              settings.gremlinPool,
              0L,
              TimeUnit.MILLISECONDS,
              queue,
              threadFactoryGremlin,
              new ThreadPoolExecutor.AbortPolicy());
    } else {
      this.gremlinExecutorService = gremlinExecutorService;
    }

    if (null == scheduledExecutorService) {
      final ThreadFactory threadFactoryGremlin = ThreadFactoryUtil.create("worker-%d");
      this.scheduledExecutorService =
          Executors.newScheduledThreadPool(settings.threadPoolWorker, threadFactoryGremlin);
    } else {
      this.scheduledExecutorService = scheduledExecutorService;
    }

    logger.info("Initialized Gremlin thread pool.  Threads in pool named with pattern gremlin-*");

    final GremlinExecutor.Builder gremlinExecutorBuilder =
        GremlinExecutor.build()
            .evaluationTimeout(settings.getEvaluationTimeout())
            .afterFailure((b, e) -> this.graphManager.rollbackAll())
            .beforeEval(b -> this.graphManager.rollbackAll())
            .afterTimeout(b -> this.graphManager.rollbackAll())
            .globalBindings(this.graphManager.getAsBindings())
            .executorService(this.gremlinExecutorService)
            .scheduledExecutorService(this.scheduledExecutorService);

    settings.scriptEngines.forEach(
        (k, v) -> {
          // use plugins if they are present
          if (!v.plugins.isEmpty()) {
            // make sure that server related classes are available at init - new approach. the
            // LifeCycleHook stuff
            // will be added explicitly via configuration using GremlinServerGremlinModule in the
            // yaml
            gremlinExecutorBuilder.addPlugins(k, v.plugins);
          }
        });

    gremlinExecutor = gremlinExecutorBuilder.create();

    logger.info("Initialized GremlinExecutor and preparing GremlinScriptEngines instances.");

    // force each scriptengine to process something so that the init scripts will fire (this is
    // necessary if
    // the GremlinExecutor is using the GremlinScriptEngineManager. this is a bit of hack, but it at
    // least allows
    // the global bindings to become available after the init scripts are run
    // (DefaultGremlinScriptEngineManager
    // runs the init scripts when the GremlinScriptEngine is created.
    settings
        .scriptEngines
        .keySet()
        .forEach(
            engineName -> {
              try {
                // use no timeout on the engine initialization - perhaps this can be a configuration
                // later
                final GremlinExecutor.LifeCycle lifeCycle =
                    GremlinExecutor.LifeCycle.build().evaluationTimeoutOverride(0L).create();
                gremlinExecutor
                    .eval("1+1", engineName, new SimpleBindings(Collections.emptyMap()), lifeCycle)
                    .join();
                registerMetrics(engineName);
                logger.info(
                    "Initialized {} GremlinScriptEngine and registered metrics", engineName);
              } catch (Exception ex) {
                logger.warn(
                    String.format(
                        "Could not initialize %s GremlinScriptEngine as init script could not be evaluated",
                        engineName),
                    ex);
              }
            });

    // script engine init may have altered the graph bindings or maybe even created new ones - need
    // to
    // re-apply those references back
    gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
        .filter(kv -> kv.getValue() instanceof Graph)
        .forEach(kv -> this.graphManager.putGraph(kv.getKey(), (Graph) kv.getValue()));

    // script engine init may have constructed the TraversalSource bindings - store them in Graphs
    // object
    gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
        .filter(kv -> kv.getValue() instanceof TraversalSource)
        .forEach(
            kv -> {
              logger.info(
                  "A {} is now bound to [{}] with {}",
                  kv.getValue().getClass().getSimpleName(),
                  kv.getKey(),
                  kv.getValue());
              this.graphManager.putTraversalSource(kv.getKey(), (TraversalSource) kv.getValue());
            });

    // determine if the initialization scripts introduced LifeCycleHook objects - if so we need to
    // gather them
    // up for execution
    hooks =
        gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
            .filter(kv -> kv.getValue() instanceof LifeCycleHook)
            .map(kv -> (LifeCycleHook) kv.getValue())
            .collect(Collectors.toList());
  }

  private void registerMetrics(final String engineName) {
    final GremlinScriptEngine engine =
        gremlinExecutor.getScriptEngineManager().getEngineByName(engineName);
    MetricManager.INSTANCE.registerGremlinScriptEngineMetrics(
        engine, engineName, "sessionless", "class-cache");
  }

  public void addHostOption(final String key, final Object value) {
    hostOptions.put(key, value);
  }

  public Map<String, Object> getHostOptions() {
    return Collections.unmodifiableMap(hostOptions);
  }

  public Object removeHostOption(final String key) {
    return hostOptions.remove(key);
  }

  public void clearHostOptions() {
    hostOptions.clear();
  }

  public ScheduledExecutorService getScheduledExecutorService() {
    return scheduledExecutorService;
  }

  public GremlinExecutor getGremlinExecutor() {
    return gremlinExecutor;
  }

  public ExecutorService getGremlinExecutorService() {
    return gremlinExecutorService;
  }

  public GraphManager getGraphManager() {
    return graphManager;
  }

  public Settings getSettings() {
    return settings;
  }

  public List<LifeCycleHook> getHooks() {
    return hooks;
  }
}
